import casadi as ca
import numpy as np

from src.utils.geometry_utils import polygon_H_from_vertices
from src.model import Model
from .convex_local_planner import ConvexLocalPlanner


class ConvexTrackingLocalPlanner(ConvexLocalPlanner):

    def __init__(self, model: Model, timestep: float, Q: np.ndarray, R: np.ndarray, N=20):
        super().__init__(model, timestep, Q, R, N)

        self._opt_u_ref = self._opti.parameter(N, self._model.num_control)
        self._opt_x_ref = self._opti.parameter(N + 1, self._model.num_state)

    def init(self):
        self._opti.subject_to(self._opt_states[0, :] == self._opt_x_ref[0, :])

        for i in range(self._N):
            state_error = self._opt_states[i, :] - self._opt_x_ref[i + 1, :]
            control_error = self._opt_controls[i, :] - self._opt_u_ref[i, :]
            self._cost += ca.mtimes([state_error, self._Q, state_error.T]) + ca.mtimes(
                [control_error, self._R, control_error.T])
        self._cost += self._slack_obs * ca.sumsqr(self._obs_slacks)
        self._opti.minimize(self._cost)

    def compute(self, next_trajectories, next_controls, obstacles):
        self._opti.set_value(self._opt_u_ref, next_controls)
        self._opti.set_value(self._opt_x_ref, next_trajectories)

        for i, obstacle in enumerate(obstacles):
            A_oi, b_oi = polygon_H_from_vertices(obstacle)
            self._opti.set_value(self._A_os[i], A_oi)
            self._opti.set_value(self._b_os[i], b_oi)

        self._opti.set_initial(self._opt_controls, self._u0)
        self._opti.set_initial(self._opt_states, self._next_states)

        sol = self._opti.solve()

        u_res = sol.value(self._opt_controls)
        next_states_pred = sol.value(self._opt_states)

        self._u0 = np.concatenate((u_res[1:], u_res[-1:]))
        self._next_states = np.concatenate((next_states_pred[1:], next_states_pred[-1:]))

        return u_res[0]
